import torch
import keras

class UltraNet_pytorch(torch.nn.Module):
    def __init__(self):
        super(UltraNet_pytorch, self).__init__()

        self.layers = torch.nn.Sequential(
            torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(16),
            torch.nn.ReLU(inplace=True),
            torch.nn.MaxPool2d(2, stride=2),

            torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(inplace=True),
            torch.nn.MaxPool2d(2, stride=2),

            torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),
            torch.nn.MaxPool2d(2, stride=2),

            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),
            torch.nn.MaxPool2d(2, stride=2),

            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),

            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),

            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),

            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU(inplace=True),

            torch.nn.Conv2d(64, 36, kernel_size=1, stride=1, padding=0)
        )
    def forward(self, x):
        x = self.layers(x)
        return x

def UltraNet_keras():
    model = keras.Sequential()

    model.add(keras.layers.Conv2D(filters=16, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False, input_shape=(160, 320, 3), name='input_layer'))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))
    model.add(keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(keras.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))
    model.add(keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))
    model.add(keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))
    model.add(keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2)))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))

    model.add(keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=(1,1), padding="same", use_bias=False))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))

    model.add(keras.layers.Conv2D(filters=36, kernel_size=(1, 1), strides=(1,1), padding="valid", name='output_layer'))

    return model
